
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#include <zlib.h>

#include "xp3.h"


static int16_t xp3_read_int16(uint8_t* buffer)
{
	return (((int64_t)buffer[1] <<  8) |
			((int64_t)buffer[0] <<  0));
}

static int32_t xp3_read_int32(uint8_t* buffer)
{
	return (((int64_t)buffer[3] << 24) |
			((int64_t)buffer[2] << 16) |
			((int64_t)buffer[1] <<  8) |
			((int64_t)buffer[0] <<  0));
}

static int64_t xp3_read_int64(uint8_t* buffer)
{
	return (((int64_t)buffer[7] << 56) |
			((int64_t)buffer[6] << 48) |
			((int64_t)buffer[5] << 40) |
			((int64_t)buffer[4] << 32) |
			((int64_t)buffer[3] << 24) |
			((int64_t)buffer[2] << 16) |
			((int64_t)buffer[1] <<  8) |
			((int64_t)buffer[0] <<  0));
}

static int64_t xp3_read_int64_from_stream(FILE* fp)
{
	uint8_t buffer[8];
	if(fread(buffer, 1, 8, fp) != 8)
	{
		return -1;
	}
	else
	{
		return xp3_read_int64(buffer);
	}
}

static wchar_t* xp3_clone_wstring(uint8_t* buffer, size_t n)
{
	wchar_t* wstring = malloc(n * sizeof(wchar_t));
	int i;
	for (i = 0; i < n; ++i)
	{
		wstring[i] = buffer[2*i] | (buffer[2*i+1] << 8);
	}
	return wstring;
}

xp3_entries_t* read_xp3_indexes(const char* xp3_path)
{
	uint8_t buffer[256];
	size_t n;
	int64_t i;
	int ei, si;
	int z;

	int64_t index_offset;
	int64_t compressed_size;
	int64_t original_size;
	uint8_t * compressed = NULL;
	uint8_t * original = NULL;

	int entry_number = 0;
	xp3_entry_t* entries = NULL;
	xp3_entries_t* index = NULL;

	uLongf uncompressed_n = 0;

	FILE* fp = fopen(xp3_path, "rb");
	if (!fp)
	{
		fprintf(stderr, "Open %s error\n", xp3_path);
		return NULL;
	}

	n = fread(buffer, 1, 11, fp);
	if (n != 11)
	{
		fclose(fp);
		fprintf(stderr, "File too small\n");
		return NULL;
	}
	if (memcmp(buffer, "XP3\r\n \n\x1a\x8bg\x01", 11))
	{
		fclose(fp);
		fprintf(stderr, "Invalid xp3 tag\n");
		return NULL;
	}
	
	index_offset = xp3_read_int64_from_stream(fp);
	if(fseek(fp, (long)index_offset, SEEK_SET))
	{
		fclose(fp);
		fprintf(stderr, "Invalid xp3 offset: %d\n", (int)index_offset);
		return NULL;
	}

	if (fgetc(fp) != 1)
	{
		fclose(fp);
		return NULL;
	}
	compressed_size = xp3_read_int64_from_stream(fp);
	original_size = xp3_read_int64_from_stream(fp);
	if (compressed_size < 0 || original_size < 0)
	{
		fclose(fp);
		return NULL;
	}

	compressed = malloc(compressed_size);
	original = malloc(original_size);
	if (!compressed || !original)
	{
		free(compressed);
		free(original);
		fclose(fp);
		fprintf(stderr, "malloc error\n");
		return NULL;
	}

	n = fread(compressed, 1, compressed_size, fp);
	if ((int64_t)n != compressed_size)
	{
		free(compressed);
		free(original);
		fclose(fp);
		fprintf(stderr, "Not enought data: %d (%d)\n", n, (int)compressed_size);
		return NULL;
	}

	uncompressed_n = (uLongf)original_size;
	z = uncompress(original, &uncompressed_n, compressed, (uLongf)compressed_size);
	if (z != Z_OK)
	{
		free(compressed);
		free(original);
		fclose(fp);
		fprintf(stderr, "zlib uncompress returns %d\n", z);
		return NULL;
	}
	if (uncompressed_n != original_size)
	{
		free(compressed);
		free(original);
		fclose(fp);
		fprintf(stderr, "uncompressed size != original size\n");
		return NULL;
	}

	for(i = 0; i < original_size;)
	{
		i += 4;
		i += xp3_read_int64(original + i) + 8;
		++entry_number;
	}
	if (i != original_size)
	{
		free(compressed);
		free(original);
		fclose(fp);
		fprintf(stderr, "Corrupted data\n");
		return NULL;
	}

	entries = malloc(entry_number*sizeof(xp3_entry_t));
	for(i = 0, ei = 0; i < original_size; ++ei)
	{
		i += 4;
		i += 8;
		i += 4;
		i += 8;
		entries[ei].encrypted = xp3_read_int32(original+i);
		i += 4;
		entries[ei].original_size = xp3_read_int64(original+i);
		i += 8;
		entries[ei].compressed_size = xp3_read_int64(original+i);
		i += 8;
		n = entries[ei].name_len = xp3_read_int16(original+i);
		i += 2;
		entries[ei].name = xp3_clone_wstring(original+i, n);
		i += n * 2;
		i += 4; // segm
		n = entries[ei].segment_number = (int)(xp3_read_int64(original+i)/28);
		i += 8;
		entries[ei].segments = malloc(n*sizeof(xp3_segment_t));
		for (si = 0; si < n; ++si)
		{
			entries[ei].segments[si].compressed = xp3_read_int32(original+i);
			i += 4;
			entries[ei].segments[si].offset = xp3_read_int64(original+i);
			i += 8;
			entries[ei].segments[si].original_size = xp3_read_int64(original+i);
			i += 8;
			entries[ei].segments[si].compressed_size = xp3_read_int64(original+i);
			i += 8;
		}
		//i += 28 * n;
		i += 4; // adlr
		i += 8;
		entries[ei].adlr = xp3_read_int32(original+i);
		i += 4;
	}

	index = malloc(sizeof(xp3_entries_t));
	index->entry_number = entry_number;
	index->entries = entries;

	free(compressed);

	if (i != original_size || ei != entry_number)
	{
		free_xp3_indexes(index);
		fclose(fp);
		fprintf(stderr, "Corrupted data\n");
		return NULL;
	}

	fclose(fp);
	return index;
}

void free_xp3_indexes(xp3_entries_t * entries)
{
	int i;
	if (entries)
	{
		for (i = 0; i < entries->entry_number; ++i)
		{
			free(entries->entries[i].name);
			free(entries->entries[i].segments);
		}
		free(entries->entries);
		free(entries);
	}
}

int decode_xp3_entry_in_buffer(const char* path, xp3_entry_t* entry, uint8_t* buffer, void (decrypter)(uint8_t* buffer, int len))
{
	int i;
	size_t n;
	int z;
	int64_t offset;
	uLongf uncompressed_n = 0;

	uint8_t* compressed_buffer = NULL;
	int compressed_buffer_size = 0;

	FILE* fp = fopen(path, "rb");
	if (!fp)
	{
		fprintf(stderr, "Open %s error\n", path);
		return -1;
	}
	
	for (i = 0, offset = 0; i < entry->segment_number; ++i)
	{
		if(fseek(fp, (long)entry->segments[i].offset, SEEK_SET))
		{
			free(compressed_buffer);
			fclose(fp);
			fprintf(stderr, "Invalid xp3 offset: %d\n", (int)entry->segments[i].offset);
			return -1;
		}

		if (!compressed_buffer || compressed_buffer_size < entry->segments[i].compressed_size)
		{
			free(compressed_buffer);
			compressed_buffer = malloc(entry->segments[i].compressed_size);
			compressed_buffer_size = (int)entry->segments[i].compressed_size;
		}
		n = fread(compressed_buffer, 1, entry->segments[i].compressed_size, fp);
		if ((int64_t)n != entry->segments[i].compressed_size)
		{
			free(compressed_buffer);
			fclose(fp);
			fprintf(stderr, "Uncompressed size incorrect\n");
			return -1;
		}

		if (entry->segments[i].compressed)
		{
			uncompressed_n = (uLongf)entry->segments[i].original_size;
			z = uncompress(buffer+offset, &uncompressed_n, compressed_buffer, (uLongf)entry->segments[i].compressed_size);
			if (z != Z_OK)
			{
				free(compressed_buffer);
				fclose(fp);
				fprintf(stderr, "zlib uncompress returns %d\n", z);
				return -1;
			}
			if (uncompressed_n != entry->segments[i].original_size)
			{
				free(compressed_buffer);
				fclose(fp);
				fprintf(stderr, "uncompressed size != original size\n");
				return -1;
			}
		}
		else
		{
			memcpy(buffer+offset, compressed_buffer, entry->segments[i].original_size);
		}
		offset += entry->segments[i].original_size;
	}
	free(compressed_buffer);
	fclose(fp);
	if (entry->original_size != offset)
	{
		fprintf(stderr, "total uncompressed size != original size\n");
		return -1;
	}
	if (decrypter) // XXX: warning C4550
	{
		decrypter(buffer, (int)entry->original_size);
	}
	else if (entry->encrypted)
	{
		fprintf(stderr, "xp3 is encrypted, but no decrypter is provided");
		return -1;
	}
	return 0;
}


